module mtblmg
    integer, parameter                          :: ng2=2, nrdir2=4
    integer, private, dimension(4) :: LTDIR_INNER ! change direction order : bottom-top to top-bottom
    data LTDIR_INNER  /1, 2, 4, 3/ 

    type mbmg
        real(8), pointer, dimension(:,:,:)      :: diagf        ! (nxy, nz, mg) diagonal block
        real(8), pointer, dimension(:,:,:,:)    :: ccrf         ! (4, nxy, nz, mg) coupling for radial
        real(8), pointer, dimension(:,:,:,:)    :: cczf
        integer, pointer, dimension(:,:,:)      :: iccr         ! (4,nxy,nz) coupling index for radial direction
        integer, pointer, dimension(:,:,:)      :: iccz         ! (2,nxy,nz) coupling index for axial direction
        integer, pointer, dimension(:)          :: nxs, nxe     ! starting and ending x-coordinates of the core region for each y-direction
        integer, pointer, dimension(:,:)        :: nodel        ! radial node index corresponding to the x- and y- cordinate
        integer :: nx, ny, nxy, nz
    end type
    interface axb1g
        module procedure axb1g2
        module procedure axb1g3
    end interface
    
contains    

subroutine newmbmg(mt, nxy, nx, ny, nz, ng)
    use allocs
    implicit none
    type(mbmg), intent(inout)   :: mt
    integer, intent(in)         :: nx, ny, nz, ng
    integer                     :: i, nxy
    ! caution
    ! nxy != nx * ny
    ! nxy means the # of real mesh  nxy <= nx * ny 
    call dmalloc(mt%diagf,nxy,nz,ng)
    call dmalloc(mt%ccrf,4,nxy,nz,ng)
    call dmalloc(mt%cczf,2,nxy,nz,ng)
    call dmalloc(mt%iccr,4,nxy,nz)
    call dmalloc(mt%iccz,2,nxy,nz)
    call dmalloc(mt%nxs,ny)
    call dmalloc(mt%nxe,ny)
    call dmalloc0(mt%nodel,0,nx+1,0,ny+1)
    mt%nx=nx
    mt%ny=ny
    mt%nxy=nxy
    mt%nz=nz
endsubroutine 

subroutine copymbmg(mt2,mt1)
    implicit none
    type(mbmg), intent(in) :: mt1
    type(mbmg), intent(out) :: mt2
    if (mt1%nz .ne. mt2%nz .or. mt1%nxy .ne. mt2%nxy) return

    mt2%diagf=mt1%diagf
    mt2%ccrf=mt1%ccrf  
    mt2%cczf=mt1%cczf  
    mt2%iccr=mt1%iccr  
    mt2%iccz=mt1%iccz  
    mt2%nodel=mt1%nodel
    mt2%nxs=mt1%nxs
    mt2%nxe=mt1%nxe
endsubroutine    
subroutine delmbmg(mt)
    implicit none
    type(mbmg), intent(inout) :: mt
    deallocate(mt%diagf)
    deallocate(mt%ccrf)
    deallocate(mt%cczf)
    deallocate(mt%nodel)
    deallocate(mt%nxs)
    deallocate(mt%nxe)
endsubroutine

! set an element
subroutine setnxse(mt, j, s, e)
    implicit none
    type(mbmg), intent(inout) :: mt
    integer, intent(in) :: j, s, e
    mt%nxs(j)=s
    mt%nxe(j)=e
endsubroutine

subroutine setnodel(mt, i, j, l)
    implicit none
    type(mbmg), intent(inout) :: mt
    integer, intent(in) :: i, j, l
    mt%nodel(i,j) =l
endsubroutine

subroutine setmbmg(mt,l,z,dir, g, value, neigh)
    implicit none
    type(mbmg), intent(inout) :: mt
    integer, intent(in) :: l        ! row in matrix
    integer, intent(in) :: z        ! index for plane
    integer, intent(in) :: dir    ! W E N S T B
    integer, intent(in) :: g        ! g11=1 g12=2 g21=3 g22=4 g1=1 g2=2
    real(8), intent(in) :: value
    integer, intent(in), optional :: neigh   
    
    if(dir .eq. 0) then
        mt%diagf(l,z,g)=value
    elseif(dir<5) then
        mt%ccrf(LTDIR_INNER(dir),l,z,g)=value
        if (PRESENT(neigh)) mt%iccr(LTDIR_INNER(dir),l,z)=neigh
    else
        mt%cczf(dir-4,l,z,g)=value
        if (PRESENT(neigh)) mt%iccz(dir-4,l,z)=neigh
    endif
endsubroutine 

! accumulate an element
subroutine accmbmg(mt,l,z,dir,g,value,neigh) 
    implicit none
    type(mbmg), intent(inout) :: mt
    integer, intent(in) :: l        ! row in matrix
    integer, intent(in) :: z        ! index for plane
    integer, intent(in) :: dir    ! N S W E T B
    integer, intent(in) :: g        ! g11=1 g12=2 g21=3 g22=4 g1=1 g2=2
    real(8), intent(in) :: value    
    integer, intent(in), optional   :: neigh    
    
    if(dir .eq. 0) then
        mt%diagf(l,z,g)=mt%diagf(l,z,g)+value
    elseif(dir<5) then
        mt%ccrf(LTDIR_INNER(dir),l,z,g)=mt%ccrf(LTDIR_INNER(dir),l,z,g)+value
        if (PRESENT(neigh)) mt%iccr(LTDIR_INNER(dir),l,z)=neigh
    else
        mt%cczf(dir-4,l,z,g)=mt%cczf(dir-4,l,z,g)+value
        if (PRESENT(neigh)) mt%iccz(dir-4,l,z)=neigh
    endif
endsubroutine 

subroutine axb1g2(mt,m,phi,aphi)
    use const
    implicit none
    type(mbmg), intent(in)       :: mt
    integer, intent(in)         :: m
    real(8), pointer            :: phi(:,:)
    real(8), pointer            :: aphi(:,:)
    integer                      :: k, l, m1, m2, idir, idirz
    real(8)                     :: aphil
    real(8), pointer, dimension(:,:,:)      :: diagf
    real(8), pointer, dimension(:,:,:,:)    :: ccrf
    real(8), pointer, dimension(:,:,:,:)    :: cczf
    integer, pointer, dimension(:,:,:)      :: iccr     
    integer, pointer, dimension(:,:,:)      :: iccz   
    
    integer :: nz, nxy  
    
    diagf=>mt%diagf;ccrf=>mt%ccrf;cczf=>mt%cczf;iccr=>mt%iccr;iccz=>mt%iccz
    nz=mt%nz;nxy=mt%nxy
    do k=1,nz
        do l=1,nxy
            aphil=diagf(l,k,m)*phi(l,k)
            do idir=1,nrdir2
                aphil=aphil+ccrf(idir,l,k,m)*phi(iccr(idir,l,k),k)
            enddo               
            aphi(l,k)=aphil
        enddo
    enddo
endsubroutine

subroutine axb1g3(mt,m,phi,aphi)
    use const
    implicit none
    type(mbmg), intent(in)       :: mt
    integer, intent(in)         :: m
    real(8), pointer            :: phi(:,:,:)
    real(8), pointer            :: aphi(:,:)
    integer                      :: k, l, m1, m2, idir, idirz
    real(8)                     :: aphil
    real(8), pointer, dimension(:,:,:)      :: diagf
    real(8), pointer, dimension(:,:,:,:)    :: ccrf
    real(8), pointer, dimension(:,:,:,:)    :: cczf
    integer, pointer, dimension(:,:,:)      :: iccr     
    integer, pointer, dimension(:,:,:)      :: iccz   
    
    integer :: nz, nxy  
    
    diagf=>mt%diagf;ccrf=>mt%ccrf;cczf=>mt%cczf;iccr=>mt%iccr;iccz=>mt%iccz
    nz=mt%nz;nxy=mt%nxy
    do k=1,nz
        do l=1,nxy
            aphil=diagf(l,k,m)*phi(l,k,m)
            do idir=1,nrdir2
                aphil=aphil+ccrf(idir,l,k,m)*phi(iccr(idir,l,k),k,m)
            enddo               
            aphi(l,k)=aphil
        enddo
    enddo
endsubroutine

endmodule